import torch
import random

def auto_mask(graph_data, node_mask_rate = 0.15, edge_mask_rate = 0.25):
    """
    对输入的图数据进行节点和边的随机剪枝（mask），
    随机将一部分节点的特征向量置为0并随机移除一部分边。

    参数:
    - graph_data: 原始的、完整的BA图数据对象。
    - node_mask_rate : 需要进行特征掩盖的节点比例。
    - edge_mask_rate : 需要移除的边的比例。

    返回:
    - masked_graph_data: 经过剪枝处理后的残缺BA图数据对象。包含了被修改的节点特征和边索引。
    - masked_info : 一个包含被剪枝信息的字典，用于后续计算loss。
      - 'masked_node_indices': 被掩盖特征的节点的索引。
      - 'original_node_features': 被掩盖节点原始的特征。
      - 'masked_edge_index': 被移除的边的索引。
      - 'original_edge_features': 被移除边的原始特征。
    """
    # 克隆数据剪枝，保留完整原始数据，用于计算loss1
    masked_graph_data = graph_data.clone()

    # 节点特征剪枝
    num_nodes = graph_data.num_nodes
    # 计算需要掩盖的节点数量
    num_masked_nodes = int(num_nodes * node_mask_rate)
    # 随机选择要掩盖的节点索引
    all_node_indices = list(range(num_nodes))
    masked_node_indices = torch.tensor(random.sample(all_node_indices, num_masked_nodes), dtype=torch.long)

    # 保存掩盖节点的原始特征作为标签，用于计算loss3
    original_node_features = masked_graph_data.x[masked_node_indices].clone()

    # 将被选中节点的特征置为0，保留拓扑结构不变
    masked_graph_data.x[masked_node_indices] = 0.0
    print(f"Masked {num_masked_nodes} node features.")

    # 边剪枝
    num_edges = graph_data.num_edges
    # 计算需要移除的边的数量
    num_masked_edges = int(num_edges * edge_mask_rate)
    # 边特征复制一份并拼接，匹配双向边的数量
    masked_graph_data.edge_attr = torch.cat([masked_graph_data.edge_attr, masked_graph_data.edge_attr], dim=0)

    # 分割出要保留的边和要移除的边
    perm = torch.randperm(num_edges)
    preserved_edge_indices = perm[num_masked_edges:]
    masked_edge_indices_in_perm = perm[:num_masked_edges]

    # 保存掩盖边的原始特征作为标签，用于计算loss2
    # original_edge_features = masked_graph_data.edge_attr[masked_edge_indices_in_perm].clone()

    # 获取被移除的边索引
    masked_edge_index = masked_graph_data.edge_index[:, masked_edge_indices_in_perm]

    # 更新图数据中的边索引为只保留的边，同时剪枝边特征
    masked_graph_data.edge_index = masked_graph_data.edge_index[:, preserved_edge_indices]
    masked_graph_data.edge_attr = masked_graph_data.edge_attr[preserved_edge_indices]
    print(f"Removed {num_masked_edges} edges.")

    masked_info = {
        'masked_node_indices': masked_node_indices,
        'original_node_features': original_node_features,  # 计算loss3
        'masked_edge_index': masked_edge_index  # 计算loss2
        # 'original_edge_features': original_edge_features
    }
    # 将masked_info附加到返回的图对象上
    masked_graph_data.masked_info = masked_info

    return masked_graph_data